import pandas as pd
import os
import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.font_manager as font_manager
from typing import List, Dict


class IndustryStockRecommender:
    def __init__(self, data_dir: str = "hs300_data"):
        """初始化荐股器，加载数据、构建图模型并配置中文字体"""
        self.data_dir = data_dir
        self.stock_df = None          # 股票核心数据
        self.stock_name_map = {}      # 股票代码→名称映射
        self.industries = []          # 所有有效行业列表
        self.G = None                 # 异构图模型（存储股票/行业关系）
        self.chinese_font = None      # 用于中文显示的字体

        # 初始化流程：字体→文件检查→数据加载→图构建
        self.chinese_font = self._get_valid_chinese_font()
        self._check_data_files()
        self._load_core_stock_data()
        self._load_stock_names()
        self._load_and_build_graph()
        
        print("✅ 初始化完成！支持行业荐股及关系图可视化")
        print("📌 提示：输入行业时可参考示例（直接复制更准确）：")
        print(f"   示例行业：{self.industries}\n")

    # ---------------- 中文字体配置 ----------------
    def _get_valid_chinese_font(self) -> str:
        """检测并返回可用的中文字体（确保图表中文显示）"""
        candidate_fonts = [
            "SimHei", "Microsoft YaHei", "PingFang SC", 
            "Heiti SC", "WenQuanYi Zen Hei"
        ]
        
        # 检测系统自带字体
        available_fonts = set(f.name for f in font_manager.fontManager.ttflist)
        for font in candidate_fonts:
            if font in available_fonts:
                print(f"✅ 检测到可用中文字体：{font}（用于图表渲染）")
                return font
        
        # 加载本地备用字体（若系统无自带字体）
        custom_font_path = os.path.join(self.data_dir, "fonts", "SimHei.ttf")
        if os.path.exists(custom_font_path):
            font_manager.fontManager.addfont(custom_font_path)
            font_manager._rebuild()
            if "SimHei" in [f.name for f in font_manager.fontManager.ttflist]:
                print(f"✅ 已加载本地字体：{custom_font_path}")
                return "SimHei"
        
        # 无可用字体时提示解决方案
        raise RuntimeError(
            "❌ 未找到可用中文字体！请按以下步骤解决：\n"
            "1. 下载 SimHei.ttf 字体文件\n"
            "2. 创建目录：hs300_data/fonts\n"
            "3. 将 SimHei.ttf 放入该目录后重试"
        )

    # ---------------- 数据检查与加载 ----------------
    def _check_data_files(self) -> None:
        """检查荐股及绘图所需的所有数据文件"""
        required_files = [
            "graph_stock_nodes.csv",    # 股票节点数据
            "graph_stock_links.csv",    # 股票联动关系
            "graph_industry_links.csv"  # 股票-行业归属关系
        ]
        for file in required_files:
            file_path = os.path.join(self.data_dir, file)
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"缺少必要文件：{file_path}\n请确保图模型数据已生成")

    def _load_core_stock_data(self):
        """加载股票核心数据（用于荐股逻辑）"""
        core_data_path = os.path.join(self.data_dir, "graph_stock_nodes.csv")
        self.stock_df = pd.read_csv(core_data_path, encoding="utf-8-sig")
        required_cols = ["stock_code", "industry", "avg_turnover_billion", "volatility"]
        missing_cols = [col for col in required_cols if col not in self.stock_df.columns]
        if missing_cols:
            raise ValueError(f"❌ 股票数据缺少必要字段：{missing_cols}")
        self.industries = sorted(self.stock_df["industry"].unique().tolist())

    def _load_stock_names(self):
        """加载股票代码-名称映射（用于显示名称）"""
        name_data_path = os.path.join(self.data_dir, "hs300_stocks_full.csv")
        if not os.path.exists(name_data_path):
            print(f"⚠️ 未找到股票名称文件{name_data_path}（将显示「名称未获取到」）")
            self.stock_name_map = {}
            return
        name_df = pd.read_csv(name_data_path, encoding="utf-8-sig")
        if "code" not in name_df.columns or "code_name" not in name_df.columns:
            print("⚠️ 股票名称文件格式不正确（缺少'code'或'code_name'列）")
            self.stock_name_map = {}
            return
        self.stock_name_map = dict(zip(name_df["code"], name_df["code_name"]))
        print(f"✅ 已加载{len(self.stock_name_map)}只股票的名称数据")

    def _load_and_build_graph(self):
        """构建异构图（包含股票节点、行业节点及关联关系）"""
        # 1. 加载图数据
        stock_nodes = pd.read_csv(
            os.path.join(self.data_dir, "graph_stock_nodes.csv"),
            encoding="utf-8-sig"
        )
        stock_links = pd.read_csv(
            os.path.join(self.data_dir, "graph_stock_links.csv"),
            encoding="utf-8-sig"
        )
        industry_links = pd.read_csv(
            os.path.join(self.data_dir, "graph_industry_links.csv"),
            encoding="utf-8-sig"
        )

        # 2. 初始化图并添加节点
        self.G = nx.Graph()
        
        # 添加股票节点（含属性：名称、行业、波动率等）
        for _, row in stock_nodes.iterrows():
            stock_code = row["stock_code"]
            self.G.add_node(
                stock_code,
                type="stock",
                stock_name=self.stock_name_map.get(stock_code, "名称未获取到"),
                industry=row["industry"],
                volatility=row["volatility"],
                avg_turnover=row["avg_turnover_billion"]
            )
        
        # 添加行业节点
        industries = industry_links["industry"].unique()
        self.G.add_nodes_from(industries, type="industry")

        # 3. 添加边（关系）
        # 股票-股票联动边（带联动系数权重）
        for _, row in stock_links.iterrows():
            self.G.add_edge(
                row["stock1"], row["stock2"],
                type="stock_corr", 
                weight=round(row["corr"], 4)  # 保留4位小数
            )
        
        # 股票-行业归属边
        for _, row in industry_links.iterrows():
            self.G.add_edge(
                row["stock_code"], row["industry"],
                type="stock_industry", 
                weight=1.0  # 归属关系权重固定为1
            )

        print(f"✅ 图模型构建完成：{len(self.G.nodes)}个节点，{len(self.G.edges)}条关系")

    # ---------------- 荐股核心逻辑 ----------------
    def _find_similar_industries(self, input_str: str) -> List[str]:
        """模糊匹配可能的行业名称"""
        return [ind for ind in self.industries if input_str in ind or ind in input_str]

    def recommend(self):
        """行业荐股主流程"""
        user_industry = input("请输入要查询的行业名称：").strip()
        if not user_industry:
            print("❌ 输入不能为空，请重新输入\n")
            return

        # 行业名称校验
        if user_industry not in self.industries:
            similar_industries = self._find_similar_industries(user_industry)
            print(f"❌ 未找到行业「{user_industry}」")
            if similar_industries:
                print(f"可能的行业名称：{similar_industries}\n")
            return

        # 筛选行业内股票并计算评分
        industry_stocks = self.stock_df[self.stock_df["industry"] == user_industry].copy()
        total_stocks = len(industry_stocks)
        if total_stocks == 0:
            print(f"❌ 行业「{user_industry}」下无股票数据\n")
            return

        # 用「成交额/波动率」作为综合评分（值越高越推荐）
        industry_stocks["score"] = industry_stocks["avg_turnover_billion"] / (industry_stocks["volatility"] + 1e-8)
        sorted_stocks = industry_stocks.sort_values(by="score", ascending=False)
        top_n = 5
        display_stocks = sorted_stocks.head(top_n)
        recommended_codes = display_stocks["stock_code"].tolist()  # 记录推荐股票代码

        # 展示推荐结果
        print(f"\n📊 「{user_industry}」行业推荐结果（共{total_stocks}只股票，展示前{len(display_stocks)}只）：")
        for idx, (_, row) in enumerate(display_stocks.iterrows(), 1):
            stock_code = row["stock_code"]
            stock_name = self.stock_name_map.get(stock_code, "【名称未获取到】")
            print(f"\n{idx}. 股票代码：{stock_code}")
            print(f"   股票名称：{stock_name}")
            print(f"   平均成交额：{round(row['avg_turnover_billion'], 2)}亿元")
            print(f"   波动率：{round(row['volatility'], 4)}")
            print(f"   推荐理由：行业内综合评分第{idx}，流动性与稳定性平衡较好")

        print("\n" + "-"*60)

        # 询问是否可视化
        visualize_choice = input("是否生成该行业与推荐股票的关系图？（输入y生成，其他键跳过）：").strip().lower()
        if visualize_choice == "y":
            self.visualize_industry_relation(industry=user_industry, stock_codes=recommended_codes)

    # ---------------- 关系图可视化 ----------------
    def visualize_industry_relation(self, industry: str, stock_codes: List[str]):
        """绘制行业与推荐股票的关系图（含股票间联动关系）"""
        # 1. 筛选可视化节点（行业+推荐股票+股票间关联节点）
        visualize_nodes = [industry] + stock_codes
        # 补充关联股票（推荐股票的联动伙伴，最多再加5个避免图过密）
        related_stocks = []
        for code in stock_codes:
            if code in self.G.nodes:
                # 取联动系数最高的2个关联股票
                neighbors = [n for n in self.G.neighbors(code) if self.G.nodes[n]["type"] == "stock"]
                neighbors_sorted = sorted(neighbors, key=lambda x: self.G[code][x]["weight"], reverse=True)
                related_stocks.extend(neighbors_sorted[:2])
        visualize_nodes += related_stocks[:5]  # 控制数量
        visualize_nodes = list(set(visualize_nodes))  # 去重
        valid_nodes = [n for n in visualize_nodes if n in self.G.nodes]
        if len(valid_nodes) < 2:
            print("❌ 可可视化节点不足，无法生成关系图")
            return

        # 2. 构建子图
        subG = self.G.subgraph(valid_nodes)

        # 3. 图表样式配置
        plt.rcParams.update({
            "font.family": self.chinese_font,
            "axes.unicode_minus": False,
            "figure.dpi": 150
        })
        fig, ax = plt.subplots(figsize=(12, 8), facecolor="white")

        # 4. 节点布局与样式
        pos = nx.spring_layout(subG, seed=42, k=0.6)  # 弹簧布局（固定seed确保布局稳定）
        node_colors = []
        node_sizes = []
        for node in subG.nodes:
            if node == industry:
                node_colors.append("#FF7F50")  # 行业节点：珊瑚色
                node_sizes.append(3000)
            elif node in stock_codes:
                node_colors.append("#4682B4")  # 推荐股票：钢青色
                node_sizes.append(1800)
            else:
                node_colors.append("#90EE90")  # 关联股票：淡绿色
                node_sizes.append(1200)

        # 5. 绘制边与节点
        nx.draw_networkx_edges(subG, pos, ax=ax, edge_color="#D3D3D3", width=1.2)
        nx.draw_networkx_nodes(
            subG, pos, ax=ax,
            node_color=node_colors, node_size=node_sizes,
            edgecolors="#333333", linewidths=1.0
        )

        # 6. 节点标签（股票显示「代码+名称」，行业显示名称）
        node_labels = {}
        for node in subG.nodes:
            if self.G.nodes[node]["type"] == "stock":
                stock_name = self.G.nodes[node]["stock_name"]
                # 名称过长时截断（避免标签重叠）
                short_name = stock_name[:6] + "..." if len(stock_name) > 6 else stock_name
                node_labels[node] = f"{node}\n{short_name}"
            else:
                node_labels[node] = node
        nx.draw_networkx_labels(
            subG, pos, ax=ax, labels=node_labels,
            font_family=self.chinese_font, font_size=9,
            font_weight="bold", font_color="#000000"
        )

        # 7. 边标签（显示股票间联动系数）
        edge_labels = {}
        for u, v, data in subG.edges(data=True):
            if data["type"] == "stock_corr" and u in stock_codes:  # 只标推荐股票的联动边
                edge_labels[(u, v)] = f"corr:{data['weight']}"
        nx.draw_networkx_edge_labels(
            subG, pos, ax=ax, edge_labels=edge_labels,
            font_family=self.chinese_font, font_size=7,
            font_color="#DC143C",  # 红色标签突出显示
            bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8)
        )

        # 8. 保存图表
        save_name = f"industry_{industry.replace('/', '_')}_relation.png"  # 替换路径特殊字符
        save_path = os.path.join(self.data_dir, save_name)
        ax.set_title(
            f"{industry}行业推荐股票关系图",
            fontdict={"fontsize": 14, "fontweight": "bold"},
            pad=20
        )
        plt.subplots_adjust(left=0.05, right=0.95, top=0.9, bottom=0.05)
        plt.savefig(save_path, dpi=300, bbox_inches="tight", facecolor="white")
        plt.close()
        print(f"✅ 关系图已保存至：{os.path.abspath(save_path)}")


# ---------------- 主程序入口 ----------------
if __name__ == "__main__":
    try:
        recommender = IndustryStockRecommender(data_dir="hs300_data")
        while True:
            recommender.recommend()
            continue_choice = input("是否继续查询其他行业？（输入y继续，其他键退出）：").strip().lower()
            if continue_choice != "y":
                print("👋 感谢使用，程序退出！")
                break
    except Exception as e:
        print(f"程序运行出错：{e}")